import cv2
import numpy as np
import torch
import torch.nn.functional as F
import torch.optim as optim
from PIL import Image
from torch.backends import cudnn
from torch.utils.data import DataLoader
from tqdm import tqdm

from Metrics import SegmentationMetric, Evaluator
from config import get_parser
from datalist import BaiDuLaneDataset
from model import DeepLab
from utils import SegmentationLosses

metric_loss = 100

'''
训练日志：
自己写的代码
目前来看训练最好的结果是miou：0.48

'''



'''

本次比赛图像分辨率非常的大，3384x1710，我的设备是两块1070Ti，可用显存15.6G，无奈只好忍痛缩小图像。下面说下主要思路：

【1】经过对label的分析，首先裁掉了图片最上部的3384x690的图像，因为这部分都是天空和树木，没有正样本的存在，裁掉后还可以减小图像的压缩比例，
一举两得。

【2】裁掉后，为了尽可能保持图像比例，在比赛不同阶段采用了3种分辨率（768x256,1024x384,1536x512）进行训练，并且大分辨率的模型是基于前一个
小分辨率的预训练模型而来（这样做的原因是大分辨率由于感受野的问题，很难训练，所以先从小分辨率入手，不仅可以快速验证模型的有效性，还能够提供较
好的特征与分布；事实证明，基于小分辨率的预训练，确实可以帮助大分辨率更好地进行收敛）。

【3】本次比赛精简了paddle提供的deeplabv3p的模型，并且根据以往在kaggle的比赛经验，重新设计了两个基于resnet残差模块的unet网络，最终成绩
了采用了三个网络的平均结果。

【4】在训练策略上，采用了修改版的Cycle LR策略（原本写了个全自动的，但是由于训练偶尔出现loss为NaN的情况，耽误了很多时间，所以只好改为手动了）。
采用Adam，前3个epochs采用默认参数训练（lr=0.001），在之后3个epochs的训练中，每个epoch平均分配出6个改变lr的地方，
改变方式为:0.001-0.0006-0.0003-0.0001-0.0004-0.0008-0.001。最后两个epoch一般采用0.0004-0.0001之间学习率训练的策略。
由于测试集与训练集在图像质量和视觉感知上差距不小，太小的学习率很容易导致过拟合，所以最小的学习率采用0.0001。
在8-10个epochs后，训练基本上就结束了。

【5】在loss function的使用上，最先的三个epochs采用sigmoid bce，后面的训练中采用bce + dice的方式，
这种方式会比单一的bce提升0.01-0.015。

【6】数据清洗上，最一开始采用了全部数据训练，发现loss经常出现不规则的跳动，经过排查，发现road 3存在几乎一半以上图像过曝的问题，
并且road 3大多在强光下拍摄，不符合测试集的分布，所以很果断的舍弃了road 3，分数也提升了0.01左右（好神奇。。。）。

【7】数据增强上，由于第五类和第八类训练数据较少（但是测试集中占比不少），所以针对这两类，采用了iaa的图像处理，从亮度、饱和度、噪点、对比度、
crop、scale等方面做了共计12000张图片的增强。最后排查了一遍road 2和road 4，把一些错误和过曝的图片都删掉了，最终保留了56000张图片
（包括数据增强的图片）进行训练。

【8】训练时采用上述分辨率，但是在实际提交结果时，添加一层bilinear，直接将结果缩放到3384x1020，实测结果会有0.005左右的提升

【9】最终的融合先使用每个模型将每个test image的结果保存为1536x512x8的npy文件，然后加载test image的三个模型的npy文件进行求平均值，
然后创建了一个bilinearNet将结果缩放成3384x1020，再将之前裁掉的3384x690的背景与结果拼接，得到最终的预测结果。

【10】【最终夺冠关键】当完成了1-9步骤后，得到的分数是0.61234分，位列第三，但是local CV的结果却不应该对应这个分数，
于是我将label叠加到了原图上，发现了一个巨大的问题，由于我训练的分辨率为1536x512，原始图像为3384x1020，我首先使用了最邻近插值来缩放label，
但是后来在生成结果时，却使用了bilinear，再加上本身分辨率并不是能够被原始分辨率整除，导致了我所有的prediction都向图像右侧偏移了4-5个像素。
于是在results_correction.py中，我将label进行了位置修正，当使用4个像素修正时，得到了当前的0.63547分。
（Note：如果使用原始图像训练、或者使用恰好被二整除关系的分辨率训练则不会出现这个问题，主要还是我设计的unet有点小问题，所以无法这样训练）

'''




class train():
    def __init__(self):
        self.args = get_parser()
        print(f"-----------{self.args.project_name}-------------")

        '''
        判断是否使用cuda
        '''
        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        '''
        设置随机数种子，保证实验结果的可重复性
        '''
        if use_cuda:
            torch.cuda.manual_seed(self.args.seed)
            torch.cuda.manual_seed_all(self.args.seed)
        else:
            torch.manual_seed(self.args.seed)
        self.device = torch.device('cuda' if use_cuda else 'cpu')
        '''
        pin_memory 锁页内存
        =True 意味着生成的Tensor数据最开始是属于内存中的锁页内存，这样将内存的Tensor转义到GPU的显存会更快一些
        主机中的内存，有两种存在方式，一是锁页，二是不锁页，
        锁页内存存放的内容在任何情况下都不会与主机的虚拟内存进行交换（注：虚拟内存就是硬盘），
        而不锁页内存在主机内存不足时，数据会存放在虚拟内存中。
        而显卡中的显存全部是锁页内存！

        当计算机的内存充足的时候，可以设置pin_memory=True。
        当系统卡住，或者交换内存使用过多的时候，设置pin_memory=False。因为pin_memory与电脑硬件性能有关，
        pytorch开发者不能确保每一个炼丹玩家都有高端设备，因此pin_memory默认为False。
        '''
        train_kwargs = {'num_workers': 0, 'pin_memory': True} if use_cuda else {}
        test_kwargs = {'num_workers': 0, 'pin_memory': True} if use_cuda else {}
        '''
        构造DataLoader
        '''
        self.train_dataset = BaiDuLaneDataset(root_file=self.args.data_base, phase='train')
        self.test_dataset = BaiDuLaneDataset(root_file=self.args.data_base, phase='test')

        #DataLoader支持多线程预读的数据加载器
        self.train_dataloader = DataLoader(self.train_dataset, batch_size=1, **train_kwargs)
        self.test_dataloader = DataLoader(self.test_dataset, batch_size=1, **test_kwargs)
        '''
        定义模型
        '''
        self.model = DeepLab().to(self.device)  # 这里可以替换模型 ，可用的有
        # deeplabv3+
        # deeplab_mobilenet

        # self.metric = SegmentationMetric(nclass=8, distributed=None)

        self.evaluator = Evaluator(8)

        '''
        CUDA加速
        '''
        if use_cuda:
            self.model = torch.nn.DataParallel(self.model, device_ids=range(torch.cuda.device_count()))
            cudnn.benchmark = True

        '''
        根据需要加载与训练模型权重参数
        '''
        if self.args.resume==True:
            data_dict = torch.load("F:/baidu_lane_seg/weights/best.pth")['model_state_dict']
            new_data_dict = {}
            for k, v in data_dict.items():
                new_data_dict[k] = v
            self.model.load_state_dict(new_data_dict, strict=True)
            print("load pretrained model successful!")
        else:
            print("initial net weights from stratch!")

        '''
        构造loss目标函数
        选择优化器
        学习率变化选择
        '''
        # 加权系数  这个的计算由来见edu
        weights = torch.FloatTensor([0.00289, 0.2411, 1.068, 2.547, 7.544, 0.2689, 0.9043, 1.572]).to(self.device)
        self.criterion = SegmentationLosses(mode='CE', weights=weights).to(self.device)
        # 这里使用了weighted crossentropy
        self.optimizer = optim.Adam(self.model.parameters(), lr=self.args.lr)
        # self.scheduler = torch.optim.lr_scheduler.MultiStepLR(self.optimizer, milestones=self.args.milestones,
        #                                                       gamma=0.1)
        self.scheduler=torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer,T_max=5,eta_min=1e-5)
        # self.scheduler=torch.optim.lr_scheduler.CyclicLR(self.optimizer,base_lr=1e-4,max_lr=5e-4)
        for epoch in range(1, self.args.epoches):
            # self.train(epoch)
            if epoch % 1 == 0:
                self.test(epoch)
        torch.cuda.empty_cache()
        print("model finish training")

    def train(self, epoch):
        global metric_loss
        self.evaluator.reset()
        self.model.train()
        average_loss = []
        pbar = tqdm(self.train_dataloader, desc=f'Train Epoch{epoch}/{self.args.epoches}')
        for data in pbar:
            img, target = data['image'], data['label']
            img, target = img.to(self.device), target.to(self.device)
            self.optimizer.zero_grad()
            outputs = self.model(img)
            loss = self.criterion(outputs, target.long()).cpu()
            loss.backward()
            self.optimizer.step()
            self.evaluator.update(target.cpu().numpy(), np.argmax(outputs.detach().cpu().numpy(), axis=1))
            average_loss.append(loss.item())
            self.scheduler.step()
            pbar.set_description(
                f'Train Epoch:{epoch}/{self.args.epoches} '
                f' train_loss:{round(np.mean(average_loss), 2)} '
                f' learning_rate:{self.optimizer.state_dict()["param_groups"][0]["lr"]}'
                f' acc:{self.evaluator.Pixel_Accuracy()}'
                f' miou:{self.evaluator.Mean_Intersection_over_Union()}')

        if np.mean(average_loss) < metric_loss and self.args.save_path:
            metric_loss = np.mean(average_loss)
            torch.save({
                'epoch': epoch,
                'model_state_dict': self.model.state_dict(),
                'optimizer_state_dict': self.optimizer.state_dict(),
                'loss': round(np.mean(average_loss), 2)
            },
                './weights/' + f'best.pth')
            print("model saved")

    def test(self, epoch):
        self.model.eval()
        self.evaluator.reset()
        with torch.no_grad():
            pbar = tqdm(self.test_dataloader, desc=f'Test Epoch{epoch}/{self.args.epoches}')
            for data in pbar:
                img, target = data['image'], data['label']
                img, target = img.to(self.device), target.to(self.device)
                outputs = self.model(img)
                preds = outputs.data.max(1)[1].cpu().numpy()
                outputs = F.interpolate(outputs, size=(1020, 3384), mode='bilinear', align_corners=True)
                self.evaluator.update(target.cpu().numpy(), np.argmax(outputs.detach().cpu().numpy(), axis=1))
                pbar.set_description(
                    f'【Test Epoch】:{epoch}/{self.args.epoches} '
                    f' acc:{self.evaluator.Pixel_Accuracy()}'
                    f' miou:{self.evaluator.Mean_Intersection_over_Union()}')

            # # 最后一个批次里的一张图拿出来看效果
            # img = img.cpu().numpy()
            #
            # img = np.transpose(img[0], axes=[1, 2, 0])
            # img *= (0.229, 0.224, 0.225)
            # img += (0.485, 0.456, 0.406)
            # img *= 255.0
            # img = img.astype(np.uint8)
            # img = cv2.resize(img, (846, 255))
            #
            # pred = preds[0].astype(np.uint8)
            # pred = self.test_dataset.decode_color_map(pred)
            #
            # img = Image.fromarray(img).convert('RGBA')
            # pred = Image.fromarray(pred).convert('RGBA')
            #
            # result = Image.blend(img, pred, 0.5)
            # result.save('./result/' + str(epoch) + '_result.png')


train = train()
